#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <libavformat/avformat.h>
#include <libavutil/log.h>

void add_udta_box(AVFormatContext *fmt_ctx, const char *imu_data) {
    // 创建 udta 盒子
    AVStream *stream = fmt_ctx->streams[0]; // 假设我们只处理第一个流
    AVPacket *pkt = av_packet_alloc();
    if (!pkt) {
        fprintf(stderr, "Could not allocate AVPacket\n");
        return;
    }

    // 创建自定义数据
    char udta_data[256];
    snprintf(udta_data, sizeof(udta_data), "IMU: %s", imu_data);

    // 将自定义数据添加到 udta 盒子
    pkt->data = (uint8_t *)udta_data;
    pkt->size = strlen(udta_data);

    // 设置时间戳
    pkt->pts = av_rescale_q(fmt_ctx->duration, AV_TIME_BASE_Q, stream->time_base);
    pkt->dts = pkt->pts; // 对于自定义数据，通常可以将 dts 设置为 pts

    // 将数据写入文件
    if (av_interleaved_write_frame(fmt_ctx, pkt) < 0) {
        fprintf(stderr, "Error writing udta box\n");
    }

    av_packet_free(&pkt);
}

int main(int argc, char *argv[]) {
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <input.mp4> <output.mp4>\n", argv[0]);
        return 1;
    }

    const char *input_file = argv[1];
    const char *output_file = argv[2];

    av_register_all();

    AVFormatContext *fmt_ctx = NULL;

    // 打开输入文件
    if (avformat_open_input(&fmt_ctx, input_file, NULL, NULL) < 0) {
        fprintf(stderr, "Could not open input file: %s\n", input_file);
        return 1;
    }

    // 查找流信息
    if (avformat_find_stream_info(fmt_ctx, NULL) < 0) {
        fprintf(stderr, "Could not find stream information\n");
        avformat_close_input(&fmt_ctx);
        return 1;
    }

    // 创建输出文件
    AVFormatContext *out_fmt_ctx = NULL;
    if (avformat_alloc_output_context2(&out_fmt_ctx, NULL, NULL, output_file) < 0) {
        fprintf(stderr, "Could not create output context\n");
        avformat_close_input(&fmt_ctx);
        return 1;
    }

    // 复制流
    for (int i = 0; i < fmt_ctx->nb_streams; i++) {
        AVStream *out_stream = avformat_new_stream(out_fmt_ctx, NULL);
        if (!out_stream) {
            fprintf(stderr, "Failed allocating output stream\n");
            return 1;
        }
        avcodec_parameters_copy(out_stream->codecpar, fmt_ctx->streams[i]->codecpar);
    }

    // 打开输出文件
    if (!(out_fmt_ctx->oformat->flags & AVFMT_NOFILE)) {
        if (avio_open(&out_fmt_ctx->pb, output_file, AVIO_FLAG_WRITE) < 0) {
            fprintf(stderr, "Could not open output file: %s\n", output_file);
            return 1;
        }
    }

    // 写入文件头
    if (avformat_write_header(out_fmt_ctx, NULL) < 0) {
        fprintf(stderr, "Error occurred when writing header\n");
        return 1;
    }

    // 添加 IMU 数据到 udta 盒子
    add_udta_box(out_fmt_ctx, "accel_x=0.1, accel_y=0.2, accel_z=0.3");

    // 写入文件尾
    av_write_trailer(out_fmt_ctx);

    // 清理
    avformat_close_input(&fmt_ctx);
    avio_closep(&out_fmt_ctx->pb);
    avformat_free_context(out_fmt_ctx);

    return 0;
}

/*
gcc -o add_udta add_udta.c -lavformat -lavcodec -lavutil
*/